#coding=utf-8
'''

    机器人数据类封装

'''
import os

from basic_math import geolib,filterlib,forcelib
from interfaces.udp import latency
import time
import numpy as np
import copy
from controller.params import Params

class Terminal(geolib.GeoEntity):
    '''
    机器人各个关键部位（例如头部，脚部等）基本描述
    其中，基类中的几何学参数(x,y,z,roll,pitch,yaw)代表
    '''
    #   静态变量（宏定义）
    defaultAccelRatio = 20.0 / 32768.0
    defaultGyroRatio = (250.0 * 3.14159265 / 180.0) / 32768.0
    def __init__(self,x=0,y=0,z=0,roll=0,pitch=0,yaw=0,cacheLen=50,name="body_part", size=[0.12,0.18,0.09],color=(0.7,0.7,0.7,1),weight=0.0,gcBias=np.array([0.0,0.0,0.0])):
        '''

        :param x: 几何学坐标
        :param y: 几何学坐标
        :param z: 几何学坐标
        :param roll: 几何学姿态
        :param pitch: 几何学姿态
        :param yaw: 几何学姿态
        :param cacheLen: 数据缓存队列
        :param name:    名称
        :param size:    几何学大小
        :param color:   颜色
        :param weight:  重量
        :param gcBias:  重心偏置（离几何学重心距离）
        '''
        super().__init__(x,y,z,roll,pitch,yaw)
        #   上次数据更新的时间戳
        self.stamp = time.time()
        #   几何学属性
        self.name=name
        self.size=size
        self.color=color
        #   重量属性
        self.weight=weight
        self.gcBias=gcBias
        #   滤波器
        self.filterLevel=2      #默认卡尔曼滤波
        self.filterKernalSize=3
        self.__filterRotation=filterlib.Kalman_Filter3D()
        #   空间描述数据
        self.__accelRaw = np.array([0.0, 0.0, 0.0])  # 加速度数据（原始）
        self.__gyroRaw = np.array([0.0, 0.0, 0.0])   # 角速度数据（原始）
        self.__postureRaw = np.array([0.0, 0.0, 0.0])   #姿态数据（原始）
        self.__positionRaw=np.array([0.0, 0.0, 0.0])    #位置数据(原始)

        #   缓存和校准数据
        self.__cacheLen=cacheLen
        self.__postureCalib=np.array([0.0, 0.0, 0.0])   #姿态校准数据
        self.__positionCalib=np.array([0.0, 0.0, 0.0])  #位置校准数据

        #   数据队列(内部私有)
        self.__postureQueue=[]
        self.__accelQueue=[]
        self.__gyroQueue=[]
        self.__gravityQueue=[]
        #   数据队列（公开）
        self.rotationQueue=[]
        self.positionQueue = []
        #   交互延迟
        self.latency=latency.Latency()

        for i in range(self.__cacheLen):
            self.__postureQueue.append([0.0,0.0,0.0])
            self.__accelQueue.append([0.0,0.0,0.0])
            self.__gyroQueue.append([0.0,0.0,0.0])
            self.__gravityQueue.append(0.0)

            self.rotationQueue.append([0.0, 0.0, 0.0])
            self.positionQueue.append([0.0, 0.0, 0.0])

        #   观测数据(经过数据过滤后)
        self.positionMonitored=np.zeros(3,float)
        self.rotationMonitored=np.zeros(3,float)
        #   模型计算数据
        self.positionKinematic = np.zeros(3, float)
        self.rotationKinematic = np.zeros(3, float)
        #   感知重力加速度
        self.G=9.8
        #   运动目标
        self.target=geolib.GeoEntity()
        #   物理学同步误差
        self.error=geolib.GeoEntity()

    def CalibrateSensor(self):
        '''
        校准传感器
        :param flag:
        :return:
        '''
        #   校准姿态和位置
        self.__postureCalib=np.mean(np.array(self.__postureQueue),axis=0)
        self.__positionCalib=np.mean(np.array(self.positionQueue),axis=0)
        #   校准重力加速度
        self.G=np.average(self.__gravityQueue)

    def UpdateFBData(self,jsonData):
        '''
        从json数据中更新基本几何学特征
        :return:
        '''
        now=time.time()
        #   线性加速度数据
        self.__accelRaw[0] = -jsonData["accel"]["x"]*Terminal.defaultAccelRatio
        self.__accelRaw[1] = -jsonData["accel"]["y"]*Terminal.defaultAccelRatio
        self.__accelRaw[2] = jsonData["accel"]["z"]*Terminal.defaultAccelRatio
        #   角速度数据
        self.__gyroRaw[0] = -jsonData["gyro"]["x"] * Terminal.defaultGyroRatio
        self.__gyroRaw[1] = -jsonData["gyro"]["y"]*Terminal.defaultGyroRatio
        self.__gyroRaw[2] = jsonData["gyro"]["z"]*Terminal.defaultGyroRatio

        #   姿态队列获取(无滤波)
        self.__postureRaw = geolib.calculate_tilt_angles(self.__accelRaw[0], self.__accelRaw[1],self.__accelRaw[2])

        self.__postureQueue.append(self.__postureRaw)
        self.__postureQueue.pop(0)
        #   加速度队列获取
        self.__accelQueue.append(self.__accelRaw)
        self.__accelQueue.pop(0)
        #   旋转队列获取
        self.__gyroQueue.append(self.__gyroRaw)
        self.__gyroQueue.pop(0)
        #   加速度队列获取
        gravity=np.linalg.norm(self.__accelRaw)
        self.__gravityQueue.append(gravity)
        self.__gravityQueue.pop(0)
        #   更新Terminal姿态
        self.UpdateGyroson()
    def UpdateGyroson(self):
        '''
        更新旋转姿态
        :return:
        '''
        now=time.time()

        #   卡尔曼滤波
        #   基于动态的G确定观测值的准确性
        rotationMonitored = self.__postureRaw - self.__postureCalib
        #   动态生成预procNoise和processNoise
        procNoise=0.00003
        observeNoise=np.abs(self.G-self.__gravityQueue[-1])*1.0+0.1
        #预测和更新
        velocity=geolib.Local2GlobalRotate(
            self.__gyroRaw[0],self.__gyroRaw[1],self.__gyroRaw[2],
            self.rotationMonitored[0],self.rotationMonitored[1],self.rotationMonitored[2])
        velocity=np.array(velocity)*(now-self.stamp)
        self.__filterRotation.Predict(velocity=velocity,procNoise=procNoise)
        rotationMonitored=self.__filterRotation.Update(observed_Pos=rotationMonitored,observeNoise=observeNoise)
        self.rotationQueue.append(rotationMonitored)
        self.rotationQueue.pop(0)
        self.rotationMonitored = self.rotationQueue[-1]
        self.stamp=now
    def GetAccelRaw(self):
        '''
        获取原始加速度和
        :return:
        '''
        return self.__accelRaw
    def GetGyroRaw(self):
        '''
        获取原始旋转速度
        :return:
        '''
        return self.__gyroRaw

    def GetRotationError(self):
        '''
        观测值误差-理论误差
        :return:
        '''
        return self.rotationMonitored[0]-self.roll,self.rotationMonitored[1]-self.pitch,self.rotationMonitored[2]-self.yaw

    def Approach(self,steps=10):
        '''
        朝向某一个Entity运动
        :param targetEntity:
        :param steps:
        :return:
        '''
        if steps==0: return
        x,y,z = self.target.DistancebyAxis(self)
        self.Translate(x/steps,y/steps,z/steps)
    def ApproachRot(self,steps=10):
        if steps==0: return
        x=self.target.roll-self.roll
        y=self.target.pitch-self.pitch
        z=self.target.yaw-self.yaw
        self.Rotate(x/steps,y/steps,z/steps)
    def ResetTarget(self):
        self.target.Mimic(self)

    def SetRefCoord(self,entity):
        '''
        设置参考坐标系（仅xyz）
        :return:
        '''
        x,y,z=-entity.VectorPose()
        self.Translate(x,y,z)
        self.target.Translate(x,y,z)

class BaseLink(Terminal):
    '''
    基座类
    '''

    def __init__(self,x,y,z,roll=0,pitch=0,yaw=0,cacheLen=30,name="",size=[0.09,0.16,0.06],color=(0.7,0.7,0.7,1),weight=0.0,gcBias=np.array([0.0,0.0,0.0])):
        '''
        初始化
        :param x:
        :param y:
        :param z:
        :param roll:
        :param pitch:
        :param yaw:
        :param cacheLen:
        :param name:
        :param size:
        :param color:
        :param weight:
        :param gcBias:
        '''
        super().__init__(x,y,z,roll,pitch,yaw,cacheLen,name,size,color,weight,gcBias)
        #   电量特征
        self.__batteryQueue=[]
        self.__powerQueue=[]
        self.__tempQueue=[]

        for i in range(cacheLen):
            self.__batteryQueue.append(0)
            self.__powerQueue.append(0)
            self.__tempQueue.append(0)

    def UpdateFBData(self, jsonData, filterKernel=5):
        '''
        从json数据包中更新数据
        :param jsonData:
        :return:
        '''
        #   更新几何学数据
        super().UpdateFBData(jsonData)
        adc=jsonData["volt"]
        self.__batteryQueue.append(adc)
        self.__batteryQueue.pop(0)
        amp = jsonData["amp"]
        self.__powerQueue.append(amp)
        self.__powerQueue.pop(0)
        self.__tempQueue.append(jsonData['temp'])
        self.__tempQueue.pop(0)
    def ADC2Volt(self,val):
        '''
        获得真实电压值
        :param val:
        :return:
        '''
        #print(val)
        return 7.2+(val-1100)/400

    def GetBattery(self):
        '''
        获得电池电量以及电压
        :return:
        '''
        volt=np.average(np.array(self.__batteryQueue))
        percentage=100*(volt-7.2)/(8.4-7.2)
        amp=self.__powerQueue[-1]
        return percentage,amp

    def GetTemperature(self):
        temp=np.average(np.array(self.__tempQueue))
        return temp

class FootTerminal(Terminal):
    '''
    脚部末端数据定义
    基于几何学类
    含有压力传感器组，识别
    '''

    def __init__(self,x,y,z,roll=0,pitch=0,yaw=0,cacheLen=30,name="",size=[0.18,0.09,0.06],color=(0.7,0.7,0.7,1),weight=0.0,gcBias=np.array([0.0,0.0,0.0]),
                 sensorPositions=[[0.04,0.025],[0.04,-0.025],[-0.04,-0.025],[-0.04,0.025]]):
        #   初始化几何学特性
        super().__init__(x,y,z,roll,pitch,yaw,cacheLen,name,size,color,weight,gcBias)
        #   压力传感器位置信息
        self.sensorPositions=sensorPositions
        self.sensorNums=len(sensorPositions)
        #   压力传感器数据
        self.__pressureRaw=np.zeros(4,float)
        #   队列构建
        self.__pressureQueue=[]     #四个压力队列
        self.pressCenterQueue=[]
        for i in range(cacheLen):
            self.__pressureQueue.append(np.ones(4))
            self.pressCenterQueue.append(np.zeros(2))
        #   开放数据
        self.pressCenter=np.zeros(2)
        self.pressure=np.zeros(4)
        self.pressureTotal=0.0
        self.supportRatio=0.0
        self.isSupport=False

    def UpdateFBData(self,jsonData):
        '''
        从json数据包中更新数据
        :param jsonData:
        :return:
        '''
        #   更新几何学数据
        super().UpdateFBData(jsonData)
        #   更新压力数据
        self.__pressureRaw[0]=jsonData["pressure"]["ADC0"]
        self.__pressureRaw[1] = jsonData["pressure"]["ADC1"]
        self.__pressureRaw[2] = jsonData["pressure"]["ADC2"]
        self.__pressureRaw[3] = jsonData["pressure"]["ADC3"]
        self.__pressureQueue.append(self.__pressureRaw)
        self.__pressureQueue.pop(0)
        self.UpdatePressureParams()
    def UpdatePressureParams(self):
        '''
        更新压力属性信息
        :return:
        '''
        #   中值滤波

        v = filterlib.median_filter_1d(self.__pressureQueue, self.filterKernalSize)
        '''
            计算过程需要进一步优化
        '''
        self.pressure=1.1 * np.power(v, -1.595)
        #self.pressure = -4.4*v+26
        self.pressureTotal = np.sum(self.pressure)
        if self.pressureTotal<Params.controlParams['terminal']['free_threshold']:
            self.pressureTotal=0.0
            self.supportRatio=0.0
            self.isSupport=False
            x,y=0,0
        else:
            self.isSupport=True
            #   计算压力中心点
            x,y=forcelib.Sensor2ZMP(self.sensorPositions,self.pressure)
            #   基于压力以及旋转，修正脚面的实际旋转值
            self.rotationMonitored[0]+=Params.controlParams['terminal']['pmod_x']*0.001*self.pressureTotal*x
            self.rotationMonitored[1] -= Params.controlParams['terminal']['pmod_y']*0.001 * self.pressureTotal * y
        self.pressCenterQueue.append(np.array([x,y]))
        self.pressCenterQueue.pop(0)
        self.pressCenter=self.pressCenterQueue[-1]


class Bone(geolib.GeoLinks):
    '''
        骨骼
        默认从Terminal出发，连接到BaseLink
        start为起点，末端为关节
    '''
    def __init__(self,startX=0,startY=0,startZ=0,length=0,name="J",weight=0.0,gcratioFromStart=0.3,size=[0.038,0.058,0,0.0],color=[0.7,0.7,0.7,1.0],jointAxis='y'):
        '''
        初始化骨骼点信息
        :param startX:
        :param startY:
        :param startZ:
        :param length:
        :param name:
        :param weight:
        :param gcratioFromStart:
        :param size:
        :param color:
        '''
        super().__init__(startX,startY,startZ,length,name)
        self.weight=weight
        self.gcRatio=gcratioFromStart
        self.size=size
        self.color=color
        self.jointAxis=jointAxis

    def GravityCenter(self):
        '''
        获取重心点
        :return:
        '''
        pos=geolib.Rotate3D(self.start.roll,self.start.pitch,self.start.yaw,0,0,self.lengh*self.gcRatio)
        return np.array([pos[0]+self.start.x,pos[1]+self.start.y,pos[2]+self.start.z])
    def SetEndRotate(self,val):
        '''
        设置关节末端的旋转
        :param val:
        :return:
        '''
        if self.jointAxis=='x':
            self.SetEndGlobalRotate(val,0.0,0.0)
        if self.jointAxis=='y':
            self.SetEndGlobalRotate(0.0,val,0.0)
        if self.jointAxis=='z':
            self.SetEndGlobalRotate(0.0, 0.0, val)

    def __str__(self):
        '''输出信息'''
        return '关节{}: 起点[{:.6f},{:.6f},{:.6f}]-->终点[{:.6f},{:.6f},{:.6f}],旋转:[{:.6f},{:.6f},{:.6f}]'\
            .format(self.name,self.start.x,self.start.y,self.start.z,self.end.x,self.end.y,self.end.z,self.start.roll,self.start.pitch,self.start.yaw)

class AppleChain:
    '''
        各个Terminal之间的骨骼链条
        解算坐标系：Global
        这个是基于AppleZhang的固定模型进行解算的，所以我们把它叫做AppleChain,这个Chain中各个肢体的长度可以调整，但是模型是固定的
        每一个从base_link到足部Terminal之间，都由一条AppleChain进行连接
    '''
    def __init__(self,baseEntity:BaseLink,terminalEntity:FootTerminal,basedockerPos=np.zeros(3),\
                 boneLength=[0.0304,0.033,0.12,0.12,0.043,0.057],\
                 boneWeights=[69.0,86.0,111.0,111.0,93.0,0.0],\
                 thresholds=[[-0.5,0.5],[-1.5,1.5],[-2.09,2.09],[-2.09,2.09],[-2.09,2.09],[-1.57,1.57]],):
        '''
        :param baseEntity: 基座对象
        :param terminalEntity: 末端对象
        初始化Chain
        :param basedockerPos:   从基座到最后一个关节的相对位置
        :param boneLength:  所有骨骼点的长度
        '''
        self.base=baseEntity
        self.terminal=terminalEntity
        self.basedockerPos=basedockerPos
        self.boneLength=boneLength
        self.boneNums=len(boneLength)
        self.thresholds=thresholds
        #   关节数值描述
        self.joints=list(np.zeros(self.boneNums))       #驱动的关节数值
        self.jointsFb=list(np.zeros(self.boneNums))     #反馈的关节数值
        self.jErrorStatic=list(np.zeros(self.boneNums)) #关节静态误差
        self.jErrorDynamic = np.zeros(self.boneNums)    #关节动态误差
        #   构建初始链
        self.chain=[]    #  关节链条
        #   重量属性
        self.massCenter=np.zeros(3)
        self.mass=np.sum(boneWeights)
        #   链属性
        #   添加第0个关节（初始关节）
        self.chain.append(Bone(self.terminal.x,self.terminal.y,self.terminal.z,
                               boneLength[0],
                               'L0',boneWeights[0],0.3,
                               [0.038,0.058,boneLength[0]],
                               jointAxis='x'))
        self.chain[0].UpdateEnd()
        for i in range(1,self.boneNums):
            self.chain.append(Bone(name="L{}".format(i),length=boneLength[i],weight=boneWeights[i],size=[0.03,0.05,boneLength[i]]))
            self.chain[-1].AppendTo(self.chain[-2])
        self.chainLen=np.sum(boneLength)
    def LoadStaticErr(self,err):
        '''
        设置静态误差
        :param err:
        :return:
        '''
        self.jErrorStatic=err
    def ResolveChain(self,autoAdjust='base',direction='z'):
        '''
        在base和terminal之间解析链条和关节
        :return:
        '''
        #   构建一个临时关节集合
        joints = np.array([0, 0, 0, 0, 0, 0], dtype=float)
        chain = copy.deepcopy(self.chain)
        chain[0].start.Mimic(self.terminal)
        chain[0].UpdateEnd()
        '''
            解析第5关节
        '''
        joints[5] = self.terminal.yaw - self.base.yaw
        chain[5].end.SetPose(geolib.Rotate3D(self.base.roll,self.base.pitch,self.base.yaw,self.basedockerPos[0],self.basedockerPos[1],self.basedockerPos[2])+self.base.VectorPose())
        chain[5].end.SetRot(self.base.VectorRot())
        chain[5].SetStartGlobalRotate(0,0,joints[5])
        chain[5].UpdateStart()
        '''
           解析第4关节 
        '''
        chain[4].end.Mimic(chain[5].start)
        #以第四关节末端为坐标系，解析它与第0关节末端的坐标特性
        x,y,z,roll,pitch,yaw=chain[4].end.LookupTransform(chain[0].end)
        joints[4]=np.arctan2(y,np.linalg.norm([z]))
        chain[4].SetStartGlobalRotate(joints[4],0,0)
        chain[4].UpdateStart()
        '''
            解析第0关节
        '''
        x, y, z, roll, pitch, yaw = chain[0].end.LookupTransform(chain[4].start)
        joints[0]=np.arctan2(y,np.linalg.norm([z]))
        chain[0].SetEndGlobalRotate(joints[0],0,0)
        chain[1].start.Mimic(chain[0].end)
        chain[1].UpdateEnd()

        '''
            解析第3,2关节
        '''

        x, y, z, roll, pitch, yaw = chain[4].start.LookupTransform(chain[1].end)
        distance=np.sqrt(x*x+z*z)
        relLen=self.chain[2].length+self.chain[3].length
        if distance>relLen:
            '''
                超出边界，需要自动调整
            '''
            relDistance=distance-relLen

            if autoAdjust=='base':
                '''
                    递归调整基座
                '''
                if direction=='xyz':
                    self.base.MoveTo(self.terminal,relDistance+0.00001)
                elif direction=='z':
                    self.base.Translate(0,0,-relDistance-0.00001)
                return self.ResolveChain(autoAdjust='base')
            if autoAdjust=='terminal':
                '''
                    递归调整Terminal
                '''
                if direction=='xyz':
                    self.terminal.MoveTo(self.base,relDistance+0.00001)
                elif direction=='z':
                    self.terminal.Translate(0,0,relDistance+0.00001)
                return self.ResolveChain(autoAdjust='terminal')
            else:
                raise Exception("{:.5f},{:.5f}".format(distance,y))
                return False
        th1,th2,th3=geolib.calculate_triangle_angles(chain[2].length,chain[3].length,distance)
        joints[3]=-(np.arctan2(x,-z)+th1)
        joints[2]=-(th3-np.pi)
        '''
            构造第2，3Link
        '''
        # print(self.chain[3])
        # print(self.chain[4])
        chain[3].end.Mimic(chain[4].start)
        chain[3].SetStartGlobalRotate(0,joints[3],0)
        chain[3].UpdateStart()

        chain[2].end.Mimic(chain[3].start)
        chain[2].SetStartGlobalRotate(0,joints[2],0)
        chain[2].UpdateStart()

        '''
            解析第一关节
        '''
        x,y,z,roll,pitch,yaw=chain[2].start.LookupTransform(chain[1].end)
        joints[1]=pitch
        '''
            判断可行性
        '''
        for i,j in enumerate(joints):
            if j>self.thresholds[i][1] or j<self.thresholds[i][0] :
                return False

        '''
            关节赋值
        '''
        self.joints=list(joints)
        self.chain=chain
        self.UpdateMassCenter()
        return True
    def BuildChain(self,autoAdjust='base'):
        '''
        基于关节以及姿态构造
        :return:
        '''
        if autoAdjust=='base':
            self.chain[0].start.Mimic(self.terminal)
            self.chain[0].UpdateEnd()
            self.chain[0].SetEndGlobalRotate(self.joints[0],0,0)
            for i in range(1,4):
                self.chain[i].AppendTo(self.chain[i-1])
                self.chain[i].SetEndGlobalRotate(0,self.joints[i],0)
            self.chain[4].AppendTo(self.chain[3])
            self.chain[4].UpdateEnd()
            self.chain[4].SetEndGlobalRotate(self.joints[4],0,0)
            self.chain[5].AppendTo(self.chain[4])
            self.chain[5].UpdateEnd()
            self.chain[5].SetEndGlobalRotate(0,0,self.joints[5])
            self.base.SetRot(self.chain[5].end.VectorRot())
            self.base.SetPose(geolib.Rotate3D(self.base.roll, self.base.pitch, self.base.yaw, -self.basedockerPos[0],
                                                 -self.basedockerPos[1], -self.basedockerPos[2]) + self.chain[5].end.VectorPose())

        elif autoAdjust=='terminal':
            self.chain[5].end.SetPose(geolib.Rotate3D(self.base.roll, self.base.pitch, self.base.yaw, self.basedockerPos[0],
                                                 self.basedockerPos[1], self.basedockerPos[2]) + self.base.VectorPose())
            self.chain[5].end.SetRot(self.base.VectorRot())
            self.chain[5].SetStartGlobalRotate(0, 0, self.joints[5])
            self.chain[5].UpdateStart()
            self.chain[4].end.Mimic(self.chain[5].start)
            self.chain[4].SetStartGlobalRotate(self.joints[4], 0, 0)
            self.chain[4].UpdateStart()
            for i in range(3,0,-1):
                self.chain[i].end.Mimic(self.chain[i+1].start)
                self.chain[i].SetStartGlobalRotate(0,self.joints[i],0)
                self.chain[i].UpdateStart()
            self.chain[0].end.Mimic(self.chain[1].start)
            self.chain[0].SetStartGlobalRotate(self.joints[0],0, 0)
            self.chain[0].UpdateStart()
            self.terminal.Mimic(self.chain[0].start)
        self.UpdateMassCenter()
    def UpdateMassCenter(self):
        '''
        获取质心
        :return:
        '''
        pAll=np.zeros(3)
        for bone in self.chain:
            p = geolib.Rotate3D(bone.start.roll, bone.start.pitch, bone.start.yaw, 0, 0, bone.length*bone.gcRatio)+bone.start.VectorPose()
            pAll+=p*bone.weight
        self.massCenter=pAll/self.mass

    def __str__(self):
        '''输出信息'''
        ret='--关节链信息--'
        i=0
        for ch in self.chain:
            ret+='\n'+str(ch)+' 关节{:3f}'.format(self.joints[i])
            i+=1
        return ret



if __name__=='__main__':
    '''
        debug
    '''
    base=BaseLink(0,-0.0,0)
    terminal=FootTerminal(0,0,0)
    chain=AppleChain(baseEntity=base,terminalEntity=terminal,basedockerPos=[0,-0.0,0],boneLength=[1,1,1,1,1,1])
    base.Mimic(chain.chain[-1].end)
    terminal.SetRot([0,0,0])
    terminal.SetPose([0,1,1])
    chain.ResolveChain()
    #print(chain)








